from fastapi import WebSocket, Query, APIRouter
import asyncio

router = APIRouter()

@router.websocket("/progress")
async def progress_websocket(websocket: WebSocket, task_id: str = Query(...)):
    """
    前端通过 ?task_id=xxx 的方式连接
    后端模拟该 task_id 的任务进度并推送
    """
    await websocket.accept()
    try:
        for i in range(101):
            # 这里你可以根据 task_id 查询任务状态
            progress_data = {
                "task_id": task_id,
                "progress": i
            }
            await websocket.send_json(progress_data)
            await asyncio.sleep(1)  # 模拟任务进度
        await websocket.close()
    except Exception as e:
        print(f"连接断开: {e}")
